import numpy as np
import platgo as pg
from scipy.spatial.distance import cdist

"""
Reference:
Q. Zhang and H. Li, MOEA/D: A multiobjective evolutionary algorithm based on decomposition,
IEEE Transactions on Evolutionary Computation, 2007, 11(6): 712-731
"""


class MOEAD(pg.Algorithm):

    type: dict = {'single': False, 'multi': True, 'many': True, 'real': True, 'binary': True, 'permutation': True,
                  "large": False, 'expensive': False, 'constrained': False, 'preference': False, 'multimodal': False,
                  'sparse': False, 'gradient': False}

    def __init__(self, maxgen: int, problem: pg.Problem) -> None:
        super().__init__(maxgen=maxgen, problem=problem)
        self.name = 'MOEA/D'
        self.xov = pg.operators.XovSbx(half=True)  # 模拟二进制交叉
        self.mut = pg.operators.MutPol(problem)  # 多项式变异
        self.decomp = pg.utils.tche  # Tchebycheff approach
        self.ps = 0.9  # probability of Select from neighbor
    
    def go(self, N: int = None, population: pg.Population = None):
        assert N or population, "N and population can't be both None"
        if population is None:
            pop = self.problem.init_pop()
        else:
            pop = population
            self.problem.N = pop.decs.shape[0]
        self.problem.cal_obj(pop)
        # TODO 种群大小可能会改变
        weight_vector, N = pg.utils.uniform_point(N=pop.N, M=self.problem.M)  # 生成权重向量, 种群大小可能有改变
        neighbor_size = max(pop.N // 10, 2)  # 确定邻域大小
        neighbor_idx = np.argsort(cdist(weight_vector, weight_vector), axis=1, kind='mergesort')[:, :neighbor_size]
        # TODO 没有考虑约束
        Z = np.nanmin(pop.objv, axis=0)  # 计算每个目标上的最小值作为理想点
        while self.not_terminal(pop):
            ps_rands = np.random.rand(pop.N)
            for i in range(pop.N):
                nei_idx = neighbor_idx[i]  # 第i个个体的邻域索引
                if ps_rands[i] < self.ps:  # 从邻域中选择父代
                    parent_idx = nei_idx[np.random.choice(neighbor_size, 2, replace=False)]
                else:  # 从整个种群不放回的选择父代
                    parent_idx = np.random.choice(pop.N, 2, replace=False)
                off = self.xov(pop=pop, p1=parent_idx[0], p2=parent_idx[1])  # 交叉
                off = self.mut(pop=off)  # 变异
                self.problem.cal_obj(off)  # 计算目标函数值
                Z = np.fmin(Z, off.objv)  # 更新理想点
                # 更新种群
                neighbor_weight = weight_vector[nei_idx]  # 邻居对应的权重向量
                nei_decomp_objv = self.decomp(pop=pop[nei_idx], weights=neighbor_weight, Z=Z)  # 计算邻居分解后的目标函数值
                off_decomp_objv = self.decomp(pop=off, weights=neighbor_weight, Z=Z)  # 计算子代在对应分解向量上的目标值
                pop[nei_idx[off_decomp_objv <= nei_decomp_objv]] = off  # 更新种群
        
        return pop
